package org.goldfish.demo.apicall.analysis.code.upload.messager.impl.tongyi;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.goldfish.demo.apicall.analysis.code.upload.messager.interfaces.LargeModelClient;
import org.goldfish.demo.apicall.analysis.code.upload.messager.interfaces.ResponseListener;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@Component
@RequiredArgsConstructor
public class ChatManager implements LargeModelClient {

    private final Config config;
    private final Generation generation = new Generation();

    @Getter
    private List<Message> userMessages = new ArrayList<>();


    @Override
    public String request(String prompt) {
        Message message = buildMessage(prompt);
        userMessages.add(message);
        log.info("");
        try {
            var response = generation.call(buildParam(false));
            return response.getOutput().getText();
        } catch (NoApiKeyException noApiKeyException) {
            throw new NullPointerException("没有api密钥" + noApiKeyException);
        } catch (InputRequiredException inputRequiredException) {
            throw new NullPointerException("请问您想问点什么" + inputRequiredException);
        }
    }

    private static Message buildMessage(String prompt) {
        return Message.builder()
                .content(prompt)
                .role(Role.USER.getValue())
                .build();
    }

    @Override
    public void requestStream(String prompt, ResponseListener listener) {
        userMessages.add(buildMessage(prompt));
        try {
            var flow = generation.streamCall(buildParam(true));
            var receiver = flow.forEach(
                    tokenRes -> {
                        log.debug("{}", tokenRes);
                        var token = tokenRes.getOutput();
                        listener.onPartialResponse(token.getText());
                    }
            );
            log.info("answer complete. result: {}", receiver);
            listener.onComplete();
        } catch (Throwable throwable) {
            listener.onError(throwable);
        }
    }

    @Override
    public void close() {
        userMessages.clear();
    }

    private GenerationParam buildParam(boolean incremental) {
        return GenerationParam.builder()
                .apiKey(config.getKey())
                .model(config.getModel())
                .messages(userMessages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(incremental)
                .build();

    }
}
